import time

import numpy as np
from datetime import datetime, timedelta
import math
import tushare as ts
import csv
pro = ts.pro_api()
class BandStrategy:
    def __init__(self, config):
        self.turnover_thresh = float(config['threshold'])
        self.increase_period = int(config['period'])

    # 筛选符合特征的个体
    def execute_strategy(self, stock_data: dict) -> dict:
        selected_stocks = {}
        # 遍历每支股票
        for stock_code in stock_data.keys():
            # 初始化区间信息
            start_index = None
            end_index = None
            continuous_increase_count = 0  # 初始化连续上涨天数计数器
            stock = stock_data[stock_code]

            # 遍历股票数据
            for date in stock.keys():
                change = stock[date]['change']

                # 判断收盘价是否上涨
                if change > 0:
                    continuous_increase_count += 1
                    if start_index is None:
                        start_index = date
                else:
                    continuous_increase_count = 0

                # 判断是否满足连续上涨天数要求
                if continuous_increase_count >= self.increase_period - 1:
                    end_index = date
                    if stock_code in selected_stocks:
                        selected_stocks[stock_code]['interval'].append((start_index, end_index))
                    else:
                        selected_stocks.setdefault(stock_code, {})
                        selected_stocks[stock_code].setdefault('interval', [])
                        selected_stocks[stock_code]['interval'].append((start_index, end_index))
        return selected_stocks

"""
筛选出具有两个涨停之间的区间的特征，1.时间长度 2. 换手率 3. 高斯加权+涨幅，并可视化这个区间之后的幅
"""
class ShortTermStrategy:
    def __init__(self, config, detect_date):
        self.config = config
        self.turnover_thresh = float(config['threshold'])
        self.increase_period = int(config['period'])
        self.gap_period_limit_up = int(config['gap_period_limit_up'])
        self.start_date = datetime.strptime(config['start_date'], '%Y%m%d')
        self.end_date = datetime.strptime(detect_date, '%Y%m%d')
        self.limit_up = 9.9

    def save_turnover(self, stock_data: dict) -> dict:
        max_min_value = {}
        turnover_csv = []
        min_turnover = float('inf')
        max_turnover = float('-inf')
        for stock_code in stock_data.keys():
            self.limit_up = 20 if stock_code.startswith("30") or stock_code.startswith("688") else 9.9
            start_index = None
            count = 0
            stock = stock_data[stock_code]
            for date in sorted(stock.keys()):
                if self.start_date <= datetime.strptime(date, "%Y%m%d") <= self.end_date:
                    if stock[date]['pct_chg'] >= self.limit_up:
                        count += 1
                        if start_index is None:
                            start_index = date
                    if count == 2:
                        end_index = date
                        start_date = datetime.strptime(start_index, '%Y%m%d')
                        end_date = datetime.strptime(end_index, '%Y%m%d')
                        interval = (end_date - start_date).days

                        if interval > self.gap_period_limit_up:
                            # 通过stock_code，date获取换手率, 换手率越高，分数越高
                            cur_date = start_date
                            turnover_rate_total = 0.0
                            while cur_date <= end_date:
                                cur_index = cur_date.strftime("%Y%m%d")
                                if cur_index in stock:
                                    turnover_rate = stock[cur_index]['turn_over']
                                    turnover_rate_total += turnover_rate
                                cur_date += timedelta(days=1)
                            min_turnover = turnover_rate_total if turnover_rate_total < min_turnover else min_turnover
                            max_turnover = turnover_rate_total if turnover_rate_total > max_turnover else max_turnover
                            turnover_csv.append((stock_code, interval, start_index, end_index, turnover_rate_total))
        field_names = ['ts_code', 'interval_days', 'start_index', 'end_index', 'turnover_rate']
        with open('./result/output.csv', 'w', newline='') as csvfile:
            writer = csv.writer(csvfile)
            # 写入列名
            writer.writerow(field_names)
            # 循环存储数据
            for row in turnover_csv:
                writer.writerow(row)
        max_min_value['max'] = max_turnover
        max_min_value['min'] = min_turnover
        return max_min_value

    def execute_strategy(self, stock_data: dict) -> dict:
        selected_stocks = {}
        if not stock_data:
            return selected_stocks
        else:
            max_min_turnover = self.save_turnover(stock_data)

            # turnover_csv = []
            for stock_code in stock_data.keys():
                self.limit_up = 20 if stock_code.startswith("30") or stock_code.startswith("688") else 9.9
                start_index = None
                count = 0
                stock = stock_data[stock_code]
                for date in sorted(stock.keys()):  # 可以指定测试的日期，新建一个sheet
                    if self.start_date <= datetime.strptime(date, "%Y%m%d") <= self.end_date:
                        if stock[date]['pct_chg'] >= self.limit_up:
                            count += 1
                            if start_index is None:
                                start_index = date

                        if count == 2:  # 符合涨停次数条件
                            second_limit = date
                            start_date = datetime.strptime(start_index, '%Y%m%d')
                            end_date = datetime.strptime(second_limit, '%Y%m%d')
                            interval = (end_date - start_date).days

                            if interval > self.gap_period_limit_up:  # 1.超参，经验->实验统计 2.多因子组合，计算框架，权重
                                # 计算得分 = 时间间隔/总时间 *1 + （区间振幅->高斯函数）*1 -> 映射,涨幅: + 换手率(归一化) // 1.区间映射表 2.周期
                                time_interval_percent = float(interval / (self.end_date - self.start_date).days)
                                amplitude = abs((stock[start_index]['high'] - stock[second_limit]['high']) /
                                                stock[start_index]['high']) * 100

                                # 通过stock_code，date获取换手率, 换手率越高，分数越高
                                cur_date = start_date
                                turnover_rate_total = 0.0
                                while cur_date <= end_date:
                                    cur_index = cur_date.strftime("%Y%m%d")
                                    if cur_index in stock:
                                        turnover_rate = stock[cur_index]['turn_over']
                                        turnover_rate_total += turnover_rate
                                    cur_date += timedelta(days=1)
                                turnover_rate_score = (turnover_rate_total - max_min_turnover['min']) / (max_min_turnover['max'] - max_min_turnover['min']) * 100
                                # score = time_interval_percent + math.exp(-amplitude) + turnover_rate_total  # 未归一化指标
                                # 判断高低位
                                # delta_price = abs(stock[second_limit]['close'] - self.sma(stock, second_limit, days=250))
                                # high_wight = self.map_delta_to_weight(delta_price, max_deviation=stock[second_limit]['close'])
                                high_wight = 1
                                new_score = high_wight * (self.get_interval_score(interval) + self.get_amplitude_score(amplitude) +
                                             turnover_rate_score)  # 归一化指标
                                if second_limit == self.end_date.strftime("%Y%m%d"):
                                    if stock_code in selected_stocks:
                                        selected_stocks[stock_code]['interval'].append((start_index, second_limit))
                                    else:
                                        selected_stocks[stock_code] = {}
                                        selected_stocks[stock_code]['interval'] = [(start_index, second_limit)]
                                        selected_stocks[stock_code]['score'] = new_score  # 最后一个区间的分数
                                        selected_stocks[stock_code]['cost_price'] = stock[second_limit]['low']
                            # 重置状态
                            count -= 1
                            start_index = date
                            second_limit = None
            return selected_stocks

    def map_delta_to_weight(self, delta_price, max_deviation):
        # 使用指数函数映射，可以根据实际需求调整指数的幂次
        exponent = -abs(delta_price) / max_deviation
        weight = math.exp(exponent)
        return weight

    def sma(self, stock_data, end_date, days=250):
        sorted_items = sorted(stock_data.items())
        # 找到end_date的索引
        i = 0
        for i, data_point in enumerate(sorted_items):
            if data_point[0] == end_date:
                break
        # 确保有足够的历史数据来计算SMA
        if i >= days - 1:
            # 计算SMA
            sma_sum = 0
            for j in range(i, i - days, -1):
                sma_sum += sorted_items[j][1]['close']
            sma_value = sma_sum / days
        else:
            # 如果历史数据不足，计算已有的i天均线水平
            sma_sum = 0
            for j in range(i, -1, -1):
                sma_sum += sorted_items[j][1]['close']
            sma_value = sma_sum / (i + 1)
        return sma_value

    @staticmethod
    def get_interval_score(time_interval):
        if time_interval < 20:
            return time_interval * 4
        elif 20 <= time_interval <= 30:
            return 80 + time_interval - 20
        elif 30 <= time_interval <= 40:
            return 90 + time_interval - 30
        else:
            return 100

    @staticmethod
    def get_amplitude_score(amplitude):
        if amplitude > 10:
            return 0
        else:
            return 100 - amplitude*10

    @staticmethod
    def get_turnover_score(turnover):
        if turnover < 7:
            return 20
        elif 7 <= turnover < 10:
            return 40
        elif 10 <= turnover < 15:
            return 60
        elif 15 <= turnover < 25:
            return 80
        else:
            return 100








